import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button
import matplotlib.patches as patches

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class ForceAnalysisArm:
    def __init__(self, l1=114, l2=60, mass=0.001):  # mass单位：kg (1g=0.001kg)
        """
        初始化双连杆机械臂受力分析模型
        :param l1: 第一连杆长度
        :param l2: 第二连杆长度
        :param mass: 末端负载质量 (kg)
        """
        self.l1 = l1
        self.l2 = l2
        self.mass = mass
        self.g = 9.81  # 重力加速度 m/s²
        self.load_force = mass * self.g  # 负载力 (N)
        
        # 初始角度 (度)
        self.theta1 = 30  # L1与x轴夹角
        self.theta2 = -30  # L2相对于L1的夹角
        
        # 工作空间限制
        self.x_min = 50
        self.x_max = 150
        self.z_max = 50  # 高度限制
        
        # 创建图形和轴
        self.fig, (self.ax1, self.ax2) = plt.subplots(1, 2, figsize=(14, 6))
        plt.subplots_adjust(left=0.1, bottom=0.4, right=0.9)  # 留出底部空间放置滑块
        
        # 设置机械臂显示轴
        max_reach = self.l1 + self.l2
        self.ax1.set_xlim(0, max_reach * 1.1)  # 只显示第一象限
        self.ax1.set_ylim(0, max_reach * 1.1)
        self.ax1.set_aspect('equal')
        self.ax1.grid(True)
        self.ax1.set_title('双连杆机械臂受力分析')
        self.ax1.set_xlabel('X坐标')
        self.ax1.set_ylabel('Z坐标 (高度)')
        
        # 绘制工作空间边界
        self.ax1.axvline(x=self.x_min, color='r', linestyle='--', alpha=0.5, label='工作空间边界')
        self.ax1.axvline(x=self.x_max, color='r', linestyle='--', alpha=0.5)
        self.ax1.axhline(y=self.z_max, color='r', linestyle='--', alpha=0.5)
        
        # 绘制底座
        self.base = plt.Circle((0, 0), 5, color='gray')
        self.ax1.add_patch(self.base)
        
        # 初始化连杆和末端执行器
        self.link1, = self.ax1.plot([], [], 'b-', lw=6)
        self.link2, = self.ax1.plot([], [], 'r-', lw=6)
        self.end_effector, = self.ax1.plot([], [], 'go', markersize=10)
        self.load_arrow = None  # 负载力箭头
        
        # 受力显示文本
        self.force_text1 = self.ax1.text(10, max_reach*0.9, '', fontsize=10)
        self.force_text2 = self.ax1.text(10, max_reach*0.8, '', fontsize=10)
        self.position_text = self.ax1.text(10, max_reach*0.7, '', fontsize=10)
        
        # 设置受力数据显示轴
        self.ax2.set_title('关节受力变化')
        self.ax2.set_xlabel('θ1角度 (°)')
        self.ax2.set_ylabel('力矩 (N·m)')
        self.ax2.grid(True)
        self.torque1_line, = self.ax2.plot([], [], 'b-', label='关节1力矩')
        self.torque2_line, = self.ax2.plot([], [], 'r-', label='关节2力矩')
        self.ax2.legend()
        
        # 存储受力数据用于绘图
        self.theta1_values = []
        self.torque1_values = []
        self.torque2_values = []
        
        # 创建滑块
        self._create_sliders()
        
        # 更新显示
        self.update_arm()
        self.update_force_plot()
        
    def _create_sliders(self):
        """创建交互滑块控件"""
        # θ1角度滑块 (度)
        ax_theta1 = plt.axes([0.1, 0.25, 0.8, 0.03])
        self.slider_theta1 = Slider(ax_theta1, 'θ1角度 (°)', 
                                  0, 90, valinit=self.theta1)  # 第一象限限制
        
        # θ2角度滑块 (度)
        ax_theta2 = plt.axes([0.1, 0.2, 0.8, 0.03])
        self.slider_theta2 = Slider(ax_theta2, 'θ2角度 (°)', 
                                  -180, 90, valinit=self.theta2)
        
        # 负载质量滑块
        ax_mass = plt.axes([0.1, 0.15, 0.8, 0.03])
        self.slider_mass = Slider(ax_mass, '负载质量 (g)', 
                                0.1, 10, valinit=self.mass*1000)
        
        # 重置按钮
        ax_reset = plt.axes([0.4, 0.05, 0.2, 0.04])
        self.button_reset = Button(ax_reset, '重置')
        
        # 绑定事件处理函数
        self.slider_theta1.on_changed(self.update_theta1)
        self.slider_theta2.on_changed(self.update_theta2)
        self.slider_mass.on_changed(self.update_mass)
        self.button_reset.on_clicked(self.reset)
        
    def calculate_position(self):
        """计算连杆端点位置"""
        # 转换角度为弧度
        theta1_rad = np.radians(self.theta1)
        theta2_rad = np.radians(self.theta2)
        total_angle = theta1_rad + theta2_rad
        
        # 计算关节点坐标
        x1 = self.l1 * np.cos(theta1_rad)
        y1 = self.l1 * np.sin(theta1_rad)
        
        # 计算末端执行器坐标
        x2 = x1 + self.l2 * np.cos(total_angle)
        y2 = y1 + self.l2 * np.sin(total_angle)
        
        return (x1, y1), (x2, y2)
    
    def calculate_forces(self):
        """计算关节受力（力矩）"""
        # 转换角度为弧度
        theta1_rad = np.radians(self.theta1)
        theta2_rad = np.radians(self.theta2)
        total_angle = theta1_rad + theta2_rad
        
        # 计算关节2的力矩 (N·m)
        # 力矩 = 力 × 力臂 × sin(角度)
        torque2 = self.load_force * self.l2 * np.sin(np.pi/2 - total_angle)
        
        # 计算关节1的力矩 (N·m)
        torque1 = self.load_force * (
            self.l1 * np.sin(np.pi/2 - theta1_rad) + 
            self.l2 * np.sin(np.pi/2 - total_angle)
        )
        
        return torque1, torque2
    
    def update_arm(self):
        """更新机械臂的显示和受力计算"""
        (x1, y1), (x2, y2) = self.calculate_position()
        torque1, torque2 = self.calculate_forces()
        
        # 检查是否在工作空间内
        in_workspace = (self.x_min <= x2 <= self.x_max) and (y2 <= self.z_max) and (x2 >= 0) and (y2 >= 0)
        
        # 更新连杆颜色以指示是否在工作空间内
        color1 = 'blue' if in_workspace else 'orange'
        color2 = 'red' if in_workspace else 'orange'
        
        # 更新连杆位置和样式
        self.link1.set_data([0, x1], [0, y1])
        self.link1.set_color(color1)
        self.link1.set_linestyle('-')
        self.link2.set_data([x1, x2], [y1, y2])
        self.link2.set_color(color2)
        self.link2.set_linestyle('-')
        self.end_effector.set_data([x2], [y2])
        
        # 绘制负载力箭头 (向下的重力)
        if self.load_arrow:
            self.load_arrow.remove()
        self.load_arrow = self.ax1.arrow(
            x2, y2, 0, -10,  # 箭头方向向下
            head_width=3, head_length=3, fc='purple', ec='purple', 
            label='负载重力' if not hasattr(self, 'arrow_label_added') else ""
        )
        if not hasattr(self, 'arrow_label_added'):
            self.ax1.legend()
            self.arrow_label_added = True
        
        # 更新受力文本
        ws_status = "在工作空间内" if in_workspace else "超出工作空间"
        self.force_text1.set_text(f'关节1力矩: {torque1:.6f} N·m')
        self.force_text2.set_text(f'关节2力矩: {torque2:.6f} N·m')
        self.position_text.set_text(f'末端位置: ({x2:.1f}, {y2:.1f}) - {ws_status}')
        
        self.fig.canvas.draw_idle()
        
        return torque1, torque2
    
    def update_force_plot(self):
        """更新受力变化曲线图"""
        # 生成theta1变化范围并计算对应力矩
        theta1_range = np.linspace(0, 90, 50)
        torque1_data = []
        torque2_data = []
        
        current_theta2 = self.theta2  # 保持theta2不变
        
        for theta1 in theta1_range:
            self.theta1 = theta1
            t1, t2 = self.calculate_forces()
            torque1_data.append(t1)
            torque2_data.append(t2)
        
        # 恢复当前theta1值
        self.theta1 = self.slider_theta1.val
        
        # 更新曲线
        self.torque1_line.set_data(theta1_range, torque1_data)
        self.torque2_line.set_data(theta1_range, torque2_data)
        
        # 调整坐标轴范围
        self.ax2.relim()
        self.ax2.autoscale_view()
        
        # 标记当前位置
        if hasattr(self, 'current_marker'):
            self.current_marker.remove()
        t1_current, _ = self.calculate_forces()
        self.current_marker = self.ax2.plot(
            [self.theta1], [t1_current], 'go', markersize=8, label='当前位置'
        )
        self.ax2.legend()
    
    # 滑块更新函数
    def update_theta1(self, val):
        self.theta1 = val
        self.update_arm()
        self.update_force_plot()
    
    def update_theta2(self, val):
        self.theta2 = val
        self.update_arm()
        self.update_force_plot()
    
    def update_mass(self, val):
        self.mass = val / 1000  # 转换为kg
        self.load_force = self.mass * self.g
        self.update_arm()
        self.update_force_plot()
    
    def reset(self, event):
        """重置所有参数到初始值"""
        self.slider_theta1.reset()
        self.slider_theta2.reset()
        self.slider_mass.reset()
        # 移除当前位置标记
        if hasattr(self, 'current_marker'):
            for marker in self.current_marker:
                marker.remove()
        self.update_force_plot()

if __name__ == "__main__":
    # 创建机械臂实例，L1=114，L2=60，末端负载1g
    arm = ForceAnalysisArm(l1=114, l2=60, mass=0.001)
    plt.show()
    